package core

import (
	"context"
	"fmt"
	"ntrace/utils/routines"
	"strings"
	"time"
)

type Aggregator struct {
	Ctx  context.Context
	Task *TraceTask
	Stat *TraceStat

	lastFlowRecord *FlowRecord
}

func NewAggregator(ctx context.Context, task *TraceTask, stat *TraceStat) *Aggregator {
	return &Aggregator{
		Ctx:  ctx,
		Task: task,
		Stat: stat,
		lastFlowRecord: &FlowRecord{
			PktRecords: make([]*PktRecord, task.MaxTTL+1),
		},
	}
}

func (a *Aggregator) StartAggregate() error {
	routines.GoWithRecover("Aggregator", a.doAggregate)
	return nil
}

func (a *Aggregator) doAggregate() {
	for {
		select {
		case <-a.Ctx.Done():
			time.Sleep(150 * time.Millisecond)
			fakeRecord := &PktRecord{}
			a.processPktRecord(a.lastFlowRecord, fakeRecord)
			return
		default:

		}

		pktRecord := a.Stat.PktRecordBucket.Peek()
		if pktRecord != nil {
			// 已经接收到响应或已超时
			delta := pktRecord.TxTs + a.Task.TimeoutMs*int64(time.Millisecond) - time.Now().UnixNano()
			if pktRecord.Received || delta <= 0 {
				pktRecord = a.Stat.PktRecordBucket.Pop()
				a.processPktRecord(a.lastFlowRecord, pktRecord)
			} else {
				delay := 100 * time.Millisecond
				if delta < int64(delay) {
					delay = time.Duration(delta)
				}
				time.Sleep(delay)
			}
		} else {
			time.Sleep(100 * time.Millisecond)
		}
	}
}

func (a *Aggregator) processPktRecord(lastFlowRecord *FlowRecord, r *PktRecord) {
	flowKey := genUnitFlowKey(a.Task.Proto, r)
	if flowKey != lastFlowRecord.Key {
		// 已收集流的所有探测
		a.processFlowRecord(lastFlowRecord)
		resetLastFlowRecord(lastFlowRecord, r, flowKey)
	}

	if r.TTL <= a.Task.MaxTTL && r.TTL > 0 {
		lastFlowRecord.PktRecords[r.TTL] = r
	}
}

func (a *Aggregator) processFlowRecord(r *FlowRecord) {
	firstReachDstTTL, lastReceiveTTL, nilCnt := a.findReceiveAndReachDstTTL(r)
	if nilCnt >= 2 {
		return
	}

	// 统计
	a.aggreFlowRecord(r, firstReachDstTTL, lastReceiveTTL)

	// 存在丢包的trace flow
	a.analyzeAbnormalFlow(r, firstReachDstTTL, lastReceiveTTL)
}

func (a *Aggregator) aggreFlowRecord(r *FlowRecord, firstReachDstTTL int, lastReceiveTTL int) {
	for ttl := int(a.Task.MaxTTL); ttl > 0; ttl-- {
		if a.Task.Proto == "tcp" && firstReachDstTTL > 0 && ttl > firstReachDstTTL {
			r.PktRecords[ttl] = nil
		}

		pktRecord := r.PktRecords[ttl]
		if pktRecord == nil {
			continue
		}

		if firstReachDstTTL > 0 {
			// 到达了目的机器
			if ttl >= firstReachDstTTL {
				a.doPktRecordStat(firstReachDstTTL, pktRecord, false, false)
			} else {
				a.doPktRecordStat(ttl, pktRecord, true, false)
			}
		} else {
			// 没到达目的机器
			if ttl > lastReceiveTTL {
				a.doPktRecordStat(lastReceiveTTL+1, pktRecord, false, true)
			} else {
				a.doPktRecordStat(ttl, pktRecord, true, false)
			}
		}
	}
}

func (a *Aggregator) analyzeAbnormalFlow(r *FlowRecord, firstReachDstTTL int, lastReceiveTTL int) {
	var text string
	var isAbnormal bool
	if firstReachDstTTL == 0 {
		// 没有到达目的地,所以是一个异常探测流
		text = analyzeUnreachDstTracePath(r, lastReceiveTTL)
		isAbnormal = true
	} else {
		text, isAbnormal = analyzeReachDstTracePath(r, firstReachDstTTL)
	}

	if isAbnormal {
		abnormal := DeepCopyFlowRecord(r)
		abnormal.Text = text
		key := genFlowKey(a.Task.Proto, r)
		abnormals := a.Stat.AbnormalFlowRecordMap[key]
		a.Stat.AbnormalFlowRecordMap[key] = append(abnormals, abnormal)
	}
}

func analyzeUnreachDstTracePath(r *FlowRecord, limitTTL int) string {
	var segs []string
	for ttl := 1; ttl <= limitTTL; ttl++ {
		pktRecord := r.PktRecords[ttl]
		if pktRecord == nil {
			segs = append(segs, "")
		} else {
			if pktRecord.Received {
				segs = append(segs, pktRecord.Device)
			} else {
				segs = append(segs, "?")
			}
		}
	}

	segs = append(segs, "******")
	return strings.Join(segs, ", ")
}

func analyzeReachDstTracePath(r *FlowRecord, limitTTL int) (string, bool) {
	var segs []string
	isAbnormal := false
	for ttl := 1; ttl <= limitTTL; ttl++ {
		pktRecord := r.PktRecords[ttl]
		if pktRecord == nil {
			segs = append(segs, "")
		} else {
			if pktRecord.Received {
				segs = append(segs, pktRecord.Device)
			} else {
				segs = append(segs, "?")
				isAbnormal = true
			}
		}
	}

	text := strings.Join(segs, ", ")
	return text, isAbnormal
}

func (a *Aggregator) findReceiveAndReachDstTTL(r *FlowRecord) (int, int, int) {
	firstReachDstTTL := 0
	lastReceiveTTL := 0
	nilCnt := 0

	// 寻找到达目标的ttl和最后一个接收到响应的探测包的ttl
	for ttl := int(a.Task.MaxTTL); ttl > 0; ttl-- {
		pktRecord := r.PktRecords[ttl]
		if pktRecord == nil {
			nilCnt++
			continue
		}

		if pktRecord.ReachDst {
			firstReachDstTTL = ttl
		}

		if pktRecord.Received && lastReceiveTTL == 0 {
			lastReceiveTTL = ttl
		}
	}

	return firstReachDstTTL, lastReceiveTTL, nilCnt
}

func (a *Aggregator) getOrCreateTraceHopStat(ttl int) *TraceHopStat {
	hopStats := a.Stat.TraceHopStats
	hopStat := hopStats[ttl]
	if hopStat == nil {
		hopStat = &TraceHopStat{
			DeviceSet: make(map[string]bool),
		}
		hopStats[ttl] = hopStat
	}

	return hopStat
}

func (a *Aggregator) doPktRecordStat(ttl int, pktRecord *PktRecord, beforeDst bool, afterReceived bool) {
	hopStat := a.getOrCreateTraceHopStat(ttl)

	if pktRecord.Device != "" {
		hopStat.DeviceSet[pktRecord.Device] = true
	} else if beforeDst {
		hopStat.DeviceSet["?"] = true
	} else if afterReceived {
		hopStat.DeviceSet["******"] = true
	}

	lastTotal := hopStat.Total
	hopStat.Total++
	if pktRecord.Received {
		latency := pktRecord.RxTs - pktRecord.TxTs
		if hopStat.MinLatency == 0 || latency < hopStat.MinLatency {
			hopStat.MinLatency = latency
		}

		if latency > hopStat.MaxLatency {
			hopStat.MaxLatency = latency
		}

		hopStat.AvgLatency = (hopStat.AvgLatency*int64(lastTotal) + latency) / int64(hopStat.Total)
	} else {
		hopStat.Loss++
	}
}

func genUnitFlowKey(proto string, r *PktRecord) string {
	var key string
	if proto == "icmp" {
		key = fmt.Sprintf("%d:%s->%s", r.Turn, r.SrcIP, r.DstIP)
	} else {
		key = fmt.Sprintf("%d:%s:%d->%s:%d", r.Turn, r.SrcIP, r.SrcPort, r.DstIP, r.DstPort)
	}
	return key
}

func genFlowKey(proto string, r *FlowRecord) string {
	var key string
	if proto == "icmp" {
		key = fmt.Sprintf("%s->%s", r.SrcIP, r.DstIP)
	} else {
		key = fmt.Sprintf("%s:%d->%s:%d", r.SrcIP, r.SrcPort, r.DstIP, r.DstPort)
	}
	return key
}

func resetLastFlowRecord(lastFlowRecord *FlowRecord, r *PktRecord, flowKey string) {
	lastFlowRecord.Key = flowKey
	lastFlowRecord.Turn = r.Turn
	lastFlowRecord.SrcIP = r.SrcIP
	lastFlowRecord.SrcPort = r.SrcPort
	lastFlowRecord.DstIP = r.DstIP
	lastFlowRecord.DstPort = r.DstPort
	for i := 0; i < len(lastFlowRecord.PktRecords); i++ {
		lastFlowRecord.PktRecords[i] = nil
	}
}
